import sys
import os
# 添加 project/ 目录到 sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import cv2
import numpy as np
import argparse
from tictoc import TicToc
import multiprocessing
from tqdm import tqdm

num_channels = 29
height = 720
width = 1280
specific_classes = [11, 1, 2, 24, 25, 27, 14, 15, 16, 17, 18, 19, 10, 9, 12, 13, 6, 7, 8, 21, 23, 20, 3, 4, 5, 26, 28,
                    0, 22]
multi_channel_array = np.zeros((num_channels, height, width))
for channel_index, value in enumerate(specific_classes):
    multi_channel_array[channel_index, :, :] = value
multi_classes_array = np.transpose(multi_channel_array, axes=(1, 2, 0))

def main(args):
    # 统计耗时
    cost = TicToc("语义图片处理")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(frames_path, args):
    camera_topics = ["camera75", "camera77", "camera80", "camera81"]
    extra_topic = "Semantic"
    output_paths = []
    output_path = os.path.join(frames_path, 'camera_epe')
    output_paths.append(output_path)
    for topic in camera_topics:
        output_path = os.path.join(frames_path, 'camera_epe', topic)
        output_paths.append(output_path)
        output_path = os.path.join(frames_path, 'camera_epe', topic, extra_topic)
        output_paths.append(output_path)
    for output_path in output_paths:
        if output_path is not None and not os.path.exists(output_path):
            print("新建目录<{}>".format(output_path))
            os.makedirs(output_path, exist_ok=True)

    for topic in camera_topics:
        path = os.path.join(frames_path, 'camera_extra', topic, extra_topic)
        file_generator = (os.path.join(root, file) for root, dirs, files in os.walk(os.path.abspath(path)) for file in
                          files)
        # 使用 tqdm 包装生成器，并设置 postfix 参数
        lines = path.strip().split('/')
        for file_path in tqdm(file_generator, postfix=os.path.join(lines[-4], lines[-2]) + "/ 语义图片处理中..."):
            # 确保文件是图片文件，这里假设你的图片文件是 .jpg 或 .png 格式
            if file_path.lower().endswith(('.jpg', '.png')):
                gt_labels = material_from_gt_label(cv2.imread(file_path))
                save_path = os.path.join(frames_path, 'camera_epe', topic, extra_topic,
                                         os.path.basename(file_path).split(".")[0] + ".npz")
                np.savez_compressed(save_path, gt_labels)


def material_from_gt_label(gt_labelmap):
    r = (multi_classes_array == gt_labelmap[:, :, 2][:, :, np.newaxis].astype(np.float32))

    class_sky = r[:, :, 0][:, :, np.newaxis]
    class_road = np.any(r[:, :, [1, 2, 3, 4, 5]], axis=2)[:, :, np.newaxis]
    class_vehicle = np.any(r[:, :, [6, 7, 8, 9, 10, 11]], axis=2)[:, :, np.newaxis]
    class_terrain = r[:, :, 12][:, :, np.newaxis]
    class_vegetation = r[:, :, 13][:, :, np.newaxis]
    class_person = np.any(r[:, :, [14, 15]], axis=2)[:, :, np.newaxis]
    class_infa = r[:, :, 16][:, :, np.newaxis]
    class_traffic_light = r[:, :, 17][:, :, np.newaxis]
    class_traffic_sign = r[:, :, 18][:, :, np.newaxis]
    class_ego = np.any(r[:, :, [19, 20]], axis=2)[:, :, np.newaxis]
    class_building = np.any(r[:, :, [21, 22, 23, 24, 25, 26]], axis=2)[:, :, np.newaxis]
    class_unlabeled = np.any(r[:, :, [27, 28]], axis=2)[:, :, np.newaxis]

    concatenated_array = np.concatenate((class_sky, class_road, class_vehicle, class_terrain, class_vegetation,
                                         class_person, class_infa, class_traffic_light, class_traffic_sign, class_ego,
                                         class_building, class_unlabeled), axis=2)
    return concatenated_array.astype(np.float32)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-path', default="/home/adt/bags/work_space/datasets",
                        help='your data root for carla')
    args = parser.parse_args()
    main(args)
